/*
 * @Author: wuqingchun
 * @Date: 2024-03-20 09:19:25
 * @Description: 
 * @LastEditors: wuqingchun
 * @LastEditTime: 2024-03-26 16:49:09
 */
#ifndef __UTCP_SOCKET_H__
#define __UTCP_SOCKET_H__

#include <stdio.h>
#include <string.h>
#include <errno.h>
#include <vector>

#include <netinet/in.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <unistd.h>
#include <linux/tcp.h>

#include "utcp_error.h"
#include "utcp_state.h"
#include "utcp_packet.h"

class CUTCPSocket : public CUTCPSocketBase
{
private:
    int32_t m_iSock {-1};
    struct sockaddr_in m_stLocalAddr {};
    struct sockaddr_in m_stPeerAddr {};
    int32_t m_iBackLog;
    CUTCPClosedState m_objCloseState;
    CUTCPOpenState m_objOpenState;
    CUTCPListenState m_objListenState;
    CUTCPSyncRecvState m_objSyncRecvState;
    CUTCPSyncSentState m_objSyncSentState;
    CUTCPEstablishedState m_objEstablishedState;
    CUTCPFinWait1State m_objFinWait1State;
    CUTCPFinWait2State m_objFinWait2State;
    CUTCPTimeWaitState m_objTimeWaitState;
    CUTCPCloseWaitState m_objCloseWaitState;
    CUTCPNullState m_objNullState;
    IUTCPState* m_pobjCurrentState {&m_objCloseState};

    std::vector<CUTCPSocket*> m_vecSubSocket;

    bool m_bDerivedSocket {false};  // 是否为派生的 socket

public:
    CUTCPSocket(bool bDerived = false):
        m_objCloseState(*this),
        m_objOpenState(*this),
        m_objListenState(*this),
        m_objSyncRecvState(*this),
        m_objSyncSentState(*this),
        m_objEstablishedState(*this),
        m_objFinWait1State(*this),
        m_objFinWait2State(*this),
        m_objTimeWaitState(*this),
        m_objCloseWaitState(*this)
    {
        m_bDerivedSocket = bDerived;
    }

    ~CUTCPSocket()
    {
        Close();
    }

    int32_t Open()
    {
        IUTCPState* pobjState = m_pobjCurrentState->Open();
        if (pobjState == nullptr)
        {
            printf("Bad state\n");
            return UTCP_FAILED;
        }
        
        m_iSock = socket(AF_INET, SOCK_DGRAM, 0);
        if (m_iSock == -1)
        {
            printf("Failed to create socket(AF_INET, SOCK_DGRAM, 0), %s\n", strerror(errno));
            return UTCP_FAILED;
        }

        m_pobjCurrentState = pobjState;

        return UTCP_SUCCESS;
    }

    int32_t Bind(const struct sockaddr *addr, socklen_t addrlen)
    {
        IUTCPState* pobjNextState = m_pobjCurrentState->Bind();
        if (pobjNextState == nullptr)
        {
            printf("Bad state\n");
            return UTCP_FAILED;
        }
        
        int32_t iErr = bind(m_iSock, addr, addrlen);
        if (iErr != 0)
        {
            printf("bind error, %s\n", strerror(errno));
            return UTCP_FAILED;
        }

        (void)memcpy(&m_stLocalAddr, addr, addrlen);
        
        if (m_stLocalAddr.sin_port == 0)
        {
            m_bDerivedSocket = true;
            // 若port指定为0,则调用bind时，系统会为其指定一个可用的端口号
            socklen_t uiAddrLen = sizeof(sockaddr_in);
            // 获取本地地址和端口
            iErr = getsockname(m_iSock, (struct sockaddr *)&m_stLocalAddr, &uiAddrLen);
            if (iErr == -1)
            {
                printf("Failed to get socket name, %s", strerror(errno));
                return UTCP_FAILED;
            }
        }
        
        m_pobjCurrentState = pobjNextState;

        return UTCP_SUCCESS;
    }

    int32_t Listen(int32_t iBackLog)
    {
        IUTCPState* pobjNextState = m_pobjCurrentState->Listen();
        if (pobjNextState == nullptr)
        {
            printf("Bad state\n");
            return UTCP_FAILED;
        }

        // if (listen(m_iSock, iBackLog) != 0)
        // {
        //     printf("listen error, %s\n", strerror(errno));
        //     return UTCP_FAILED;
        // }

        m_iBackLog = iBackLog;

        m_pobjCurrentState = pobjNextState;

        return UTCP_SUCCESS;
    }

    CUTCPSocket* Accept(struct sockaddr *addr, socklen_t *addrlen)
    {
        IUTCPState* pobjNextState = m_pobjCurrentState->Accept();
        if (pobjNextState == nullptr)
        {
            printf("Bad state\n");
            return nullptr;
        }

        if (__RecvSync((sockaddr_in*)addr, addrlen) != UTCP_SUCCESS)
        {
            printf("bad message\n");
            return nullptr;
        }

        return __DeriveSocket(addr, *addrlen);
    }

    int32_t Connect(const struct sockaddr *pstPeerAddr, socklen_t uiAddrLen)
    {
        m_stPeerAddr = *(sockaddr_in *)pstPeerAddr;

        if (m_bDerivedSocket)
        {
            return __DoPassiveOpen();
        }   

        return __DoActiveOpen();
    }

    ssize_t Send(const void* pBuf, size_t ulBufLen, int iFlags)
    {
        //return m_pobjCurrentState->Send();
        char acBuf[1024];

        tcphdr *pstHdr = (tcphdr *)acBuf;
        pstHdr->psh = 1;

        memcpy((void*)(pstHdr + 1), pBuf, ulBufLen);

        return sendto(m_iSock, acBuf, sizeof(tcphdr) + ulBufLen, 0, (const sockaddr*)&m_stPeerAddr, sizeof(m_stPeerAddr));
    }

    ssize_t Recv(void* pBuf, size_t ulBufLen, int iFlags)
    {
        char acBuf[1024];

        sockaddr_in stPeerAddr = {};
        socklen_t uiAddLen = sizeof(sockaddr_in);

        ssize_t lRecv = recvfrom(m_iSock, acBuf, 1024, 0, (sockaddr *)&stPeerAddr, &uiAddLen);
        if (lRecv <= 0)
        {
            printf("recvfrom error, %s\n", strerror(errno));
            return -1;
        }

        if (!__IsSameAddr(&m_stPeerAddr, &stPeerAddr))
        {
            return -1;
        }

        memcpy(pBuf, acBuf + sizeof(tcphdr), lRecv - sizeof(tcphdr));

        return lRecv - sizeof(tcphdr);
    }

    void Close()
    {
        if (m_iSock != -1)
        {
            close(m_iSock);
            m_iSock = -1;
        }
        
        m_pobjCurrentState->Close();
    }

    IUTCPState& GetClosedState()
    {
        return m_objCloseState;
    }

    IUTCPState& GetListenState()
    {
        return m_objListenState;
    }


    IUTCPState& GetSyncSentState()
    {
        return m_objSyncSentState;
    }

    IUTCPState& GetSyncRecvState()
    {
        return m_objSyncRecvState;
    }
    
    IUTCPState& GetEstablishedState()
    {
        return m_objSyncRecvState;
    }

    IUTCPState& GetFinWait1State()
    {
        return m_objFinWait2State;
    }

    IUTCPState& GetTimeWaitState()
    {
        return m_objTimeWaitState;
    }
    
    IUTCPState& GetCloseWaitState()
    {
        return m_objCloseWaitState;
    }

    IUTCPState& GetNullState()
    {
        return m_objNullState;
    }
private:

    int32_t __CreateSocket()
    {
        int32_t iSock = socket(AF_INET, SOCK_DGRAM, 0);
        if (m_iSock == -1)
        {
            printf("Failed to create socket(AF_INET, SOCK_DGRAM, 0), %s\n", strerror(errno));
            return UTCP_FAILED;
        }

        int32_t iOn = 1;
        int32_t iErr = setsockopt(iSock, SOL_SOCKET, SO_REUSEADDR, &iOn, sizeof(iOn));
        if (iErr == -1)
        {
            close(iSock);
            printf("setsockopt error, %s\n", strerror(errno));
            return UTCP_FAILED;
        }

        return iSock;
    }

    CUTCPSocket* __DeriveSocket(const struct sockaddr *pstPeerAddr, socklen_t uiAddrLen)
    {
        // 派生子socket
        CUTCPSocket* pobjSubSocket = new (std::nothrow) CUTCPSocket(true);
        if (pobjSubSocket == nullptr)
        {
            delete pobjSubSocket;
            return nullptr;
        }

        if (pobjSubSocket->Open() != UTCP_SUCCESS)
        {
            delete pobjSubSocket;
            return nullptr;
        }

        sockaddr_in stLocal = {};
        stLocal.sin_family = AF_INET;
        if (pobjSubSocket->Bind((const struct sockaddr *)&stLocal, sizeof(stLocal)) != UTCP_SUCCESS)
        {
            delete pobjSubSocket;
            return nullptr;
        }

        if (pobjSubSocket->Connect(pstPeerAddr, uiAddrLen) != UTCP_SUCCESS)
        {
            delete pobjSubSocket;
            return nullptr;
        }
    
        m_vecSubSocket.push_back(pobjSubSocket);

        return pobjSubSocket;
    }

    ssize_t __SendSync()
    {
        tcphdr stHdr = {};
        stHdr.syn = 1;
        
        printf("%s\n", __FUNCTION__);
        return sendto(m_iSock, &stHdr, sizeof(stHdr), 0, (const sockaddr *)&m_stPeerAddr, sizeof(m_stPeerAddr));
    }

    ssize_t __SendAck()
    {
        tcphdr stHdr = {};
        stHdr.ack = 1;

        printf("%s\n", __FUNCTION__);
        return sendto(m_iSock, &stHdr, sizeof(stHdr), 0, (const sockaddr *)&m_stPeerAddr, sizeof(m_stPeerAddr));
    }

    ssize_t __SendSyncAck()
    {
        tcphdr stHdr = {};
        stHdr.syn = 1;
        stHdr.ack = 1;

        printf("%s\n", __FUNCTION__);
        return sendto(m_iSock, &stHdr, sizeof(stHdr), 0, (const sockaddr *)&m_stPeerAddr, sizeof(m_stPeerAddr));
    }

    int32_t __RecvSyncAck(sockaddr_in* pstPeerAddr, socklen_t* puiAddLen)
    {
        tcphdr stHdr;

        printf("%s\n", __FUNCTION__);
        ssize_t lRecv = recvfrom(m_iSock, &stHdr, sizeof(stHdr), 0, (sockaddr *)pstPeerAddr, puiAddLen);
        if (lRecv <= 0)
        {
            printf("recvfrom error, %s\n", strerror(errno));
            return UTCP_FAILED;
        }

        if (stHdr.syn == 0 || stHdr.ack == 0)
        {
            return UTCP_FAILED;
        }

        return UTCP_SUCCESS;
    }

    int32_t __RecvSync(sockaddr_in* pstPeerAddr, socklen_t* puiAddLen)
    {
        tcphdr stHdr;

        printf("%s\n", __FUNCTION__);

        ssize_t lRecv = recvfrom(m_iSock, &stHdr, sizeof(stHdr), 0, (sockaddr *)pstPeerAddr, puiAddLen);
        if (lRecv <= 0)
        {
            printf("recvfrom error, %s\n", strerror(errno));
            return UTCP_FAILED;
        }

        if (!stHdr.syn)
        {
            return UTCP_FAILED;
        }

        return UTCP_SUCCESS;
    }

    int32_t __RecvAck(sockaddr_in* pstPeerAddr, socklen_t* puiAddLen)
    {
        tcphdr stHdr = {};
        
        printf("%s\n", __FUNCTION__);

        ssize_t lRecv = recvfrom(m_iSock, &stHdr, sizeof(stHdr), 0, (sockaddr *)pstPeerAddr, puiAddLen);
        if (lRecv <= 0)
        {
            printf("recvfrom error, %s\n", strerror(errno));
            return UTCP_FAILED;
        }

        if (!stHdr.ack)
        {
            return UTCP_FAILED;
        }

        return UTCP_SUCCESS;
    }

    /**
     * @brief 处理主动打开连接
     * 
     * @return int32_t 
     */
    int32_t __DoActiveOpen()
    {
        ssize_t lSendRet;
        sockaddr_in stPeerAddr = {};
        socklen_t uiAddLen = sizeof(sockaddr_in);

        lSendRet = __SendSync();
        if (lSendRet <= 0)
        {
            printf("__SendSync error, %s\n", strerror(errno));
            return UTCP_FAILED;
        }

        // 此处更换地址
        if (__RecvSyncAck(&stPeerAddr, &uiAddLen) != UTCP_SUCCESS)
        {
            printf("__RecvSyncAck error\n");
            return UTCP_FAILED;
        }

        m_stPeerAddr = stPeerAddr;

        // 检查对端地址
        // if (!__IsSameAddr(&m_stPeerAddr, &stPeerAddr))
        // {
        //     printf("__IsSameAddr error\n");
        //     return UTCP_FAILED;
        // }

        lSendRet = __SendAck();
        if (lSendRet <= 0)
        {
            printf("__SendAck error, %s\n", strerror(errno));
            return UTCP_FAILED;
        }

        return UTCP_SUCCESS;
    }

    /**
     * @brief 处理被动打卡连接
     * 
     * @return int32_t 
     */
    int32_t __DoPassiveOpen()
    {
        ssize_t lSendRet;
        sockaddr_in stPeerAddr = {};
        socklen_t uiAddLen = sizeof(sockaddr_in);

        lSendRet = __SendSyncAck();
        if (lSendRet <= 0)
        {
            printf("__SendSyncAck error, %s\n", strerror(errno));
            return UTCP_FAILED;
        }

        if (__RecvAck(&stPeerAddr, &uiAddLen) != UTCP_SUCCESS)
        {
            return UTCP_FAILED;
        }

        // 检查对端地址
        if (!__IsSameAddr(&m_stPeerAddr, &stPeerAddr))
        {
            return UTCP_FAILED;
        }

        return UTCP_SUCCESS;
    }

    bool __IsSameAddr(const sockaddr_in* pstSrc, const sockaddr_in* pstDst)
    {
        if (pstSrc->sin_addr.s_addr == pstDst->sin_addr.s_addr &&
            pstSrc->sin_family == pstDst->sin_family &&
            pstSrc->sin_port == pstDst->sin_port)
        {
            return true;
        }

        return false;
    }
};

#endif //__UTCP_SOCKET_H__